/* gibbsDemoDenoising
 * Denoising of letter A using Gibbs sampling
 * with an Ising Prior and a Gaussian likelihood
 * Code written by Aurelien Lucchi
 * Based on matlab code originally written by Brani Vidakovic and modified by Kevin Murphy
 */

#include <cv.h>
#include <highgui.h>
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

int main(int argc, char *argv[])
{
  const int nLabels = 2;
  const double sigma = 0.5; // noise level
  const int J = 1;
  long n,i;

  int maxIter = 1e7;
  if(argc > 1) {
    maxIter = atof(argv[1]);
  }
  printf("Maximum number of iterations = %d\n", maxIter);

  gsl_rng *r = gsl_rng_alloc(gsl_rng_mt19937);

  double mus[nLabels];
  mus[0] = -1;
  mus[1] = 1;

  // load input matrix consisting of letter A. The body of letter
  // A is made of 1's while the background is made of -1's.
  IplImage* img = cvLoadImage("lettera.bmp");
  CvScalar cv_avg = cvAvg(img);
  double avg = cv_avg.val[0];
  printf("Average pixel value = %g\n", avg);

  // allocate memory
  ulong nPixels = img->width*img->height;
  double* img_n = new double[nPixels];
  double** local_evidence = new double*[nPixels];
  for(int n = 0; n < nPixels; ++n) {
    local_evidence[n] = new double[nLabels];
  }
  double* inferred_labels = new double[nPixels];

  // create noisy image
  uchar* pData;
  uchar* pData2;
  IplImage* img2 = cvCreateImage(cvSize(img->width, img->height), IPL_DEPTH_8U, 1);
  int idx = 0;
  for(int h = 0; h < img->height; ++h) {
    for(int w = 0; w < img->width; ++w) {
      pData = (((uchar*)(img->imageData + h*img->widthStep) + w*img->nChannels));
      pData2 = (((uchar*)(img2->imageData + h*img2->widthStep) + w*img2->nChannels));
      if(*pData < avg) {
        // black pixel => class 0
        img_n[idx] = mus[0];
      } else {
        img_n[idx] = mus[1];
      }
      img_n[idx] += gsl_ran_gaussian(r,sigma);

      // only for display
      *pData2 = img_n[idx]*50+100;

      // get local evidence
      for(int l = 0; l < nLabels; ++l) {
        local_evidence[idx][l] = gsl_ran_gaussian_pdf(img_n[idx]-mus[l],sigma);
      }

      ++idx;
    }
  }

  cvSaveImage("init.png", img2);

  idx = 0;
  for(int h = 0; h < img->height; ++h) {
    for(int w = 0; w < img->width; ++w) {
      pData2 = (((uchar*)(img2->imageData + h*img2->widthStep) + w*img2->nChannels));

      double max_evidence = local_evidence[idx][0];
      int max_label = 0;
      for(int l = 1; l < nLabels; ++l) {
        if(local_evidence[idx][l] > max_evidence) {
          max_evidence = local_evidence[idx][l];
          max_label = l;
        }
      }

      *pData2 = max_label*255.0/(nLabels-1);
      ++idx;
    }
  }
  cvSaveImage("max_evidence_0.png", img2);

  // Start gibbs sampling
  int ridx;
  for(int i = 0; i < maxIter; ++i) {
    // Select a pixel at random
    int w = gsl_rng_uniform_int(r, img->width);
    int h = gsl_rng_uniform_int(r, img->height);

    // accumlate weights for neighbors
    double wi = 0;
    if(w > 0) {
      ridx = (h*img->width) + (w-1);
      wi += inferred_labels[ridx];
    }
    if(w < (img->width-1)) {
      ridx = (h*img->width) + (w+1);
      wi += inferred_labels[ridx];
    }
    if(h > 0) {
      ridx = ((h-1)*img->width) + w;
      wi += inferred_labels[ridx];
    }
    if(h < (img->height-1)) {
      ridx = ((h+1)*img->width) + w;
      wi += inferred_labels[ridx];
    }

    ridx = (h*img->width) + w;
    // Compute probability p(x_i) = exp(-E(x_i))
    // E(x_i) = -(\sum_j J*x_i*x_j + h(x_i))
    // where h(x_i) is the local evidence
    double p1 = exp(J*wi) * local_evidence[ridx][1];  //xi=+1 => +J in exp
    double p0 = exp(-J*wi) * local_evidence[ridx][0]; //xi=-1 => -J in exp
    double prob = p1/(p0+p1);
    if(gsl_rng_uniform(r) < prob) {
      inferred_labels[ridx] = 1;
    } else {
      inferred_labels[ridx] = -1;
    }
  }

  // output inferred image
  idx = 0;
  for(int h = 0; h < img->height; ++h) {
    for(int w = 0; w < img->width; ++w) {
      pData2 = (((uchar*)(img2->imageData + h*img2->widthStep) + w*img2->nChannels));
      int label = (inferred_labels[idx]+1)/2;
      *pData2 = label*255.0/(nLabels-1);
      ++idx;
    }
  }
  cvSaveImage("evidences.png", img2);

  // Cleaning
  for(int n = 0; n < nPixels; ++n) {
    delete[] local_evidence[n];
  }
  delete[] local_evidence;
  delete[] inferred_labels;
  delete[] img_n;
  cvReleaseImage(&img);
  cvReleaseImage(&img2);

  return 0;
}
